//
// Created by wchenbt on 3/4/2021.
//
#ifdef ENABLE_MAPLE_SAN

#include "asan_stackvar.h"

#include <string_utils.h>

#include "san_common.h"
#include <set>

#define ENABLE_STACK_SIZE_LIMIT false

namespace maple {

// Fake stack allocator (asan_fake_stack.h) has 11 size classes
// for every power of 2 from kMinStackMallocSize to kMaxAsanStackMallocSizeClass
static int StackMallocSizeClass(uint64_t LocalStackSize) {
  // unblock the stack size limit
  if (ENABLE_STACK_SIZE_LIMIT) {
    CHECK_FATAL(LocalStackSize <= kMaxStackMallocSize, "Too large stack size.");
  }
  uint64_t MaxSize = kMinStackMallocSize;
  for (int i = 0;; i++, MaxSize *= 2) {
    if (LocalStackSize <= MaxSize) {
      return i;
    }
  }
}

FunctionStackPoisoner::FunctionStackPoisoner(MeFunction &function, AddressSanitizer &asan)
    : ASan(asan),
      meFunction(&function),
      mirFunction(function.GetMirFunc()),
      module(mirFunction->GetModule()),
      IntptrTy(asan.IntPtrTy),
      Mapping(asan.Mapping),
      StackAlignment(1 << Mapping.Scale) {
  IntptrPtrTy = GlobalTables::GetTypeTable().GetOrCreatePointerType(*IntptrTy, PTY_ptr);
  // initialize function arguments symbols
  for (size_t i = 0; i < mirFunction->GetFormalCount(); ++i) {
    const MIRSymbol *symbolPtr = mirFunction->GetFormal(i);
    callArgSymbols.insert(symbolPtr);
    callArgSymbolNames.insert(symbolPtr->GetName());
  }
  initializeCallbacks(*module);
}

void FunctionStackPoisoner::copyToShadow(const std::vector<uint8_t> ShadowMask, const std::vector<uint8_t> ShadowBytes,
                                         MIRBuilder *mirBuilder, BaseNode *ShadowBase, StmtNode *InsBefore) {
  copyToShadow(ShadowMask, ShadowBytes, 0, ShadowMask.size(), mirBuilder, ShadowBase, InsBefore);
}

void FunctionStackPoisoner::copyToShadow(const std::vector<uint8_t> ShadowMask, const std::vector<uint8_t> ShadowBytes,
                                         size_t Begin, size_t End, MIRBuilder *mirBuilder, BaseNode *ShadowBase,
                                         StmtNode *InsBefore) {
  assert(ShadowMask.size() == ShadowBytes.size());
  size_t Done = Begin;
  const size_t threshold = 64;
  for (size_t i = Begin, j = Begin + 1; i < End; i = j++) {
    if (!ShadowMask[i]) {
      assert(!ShadowBytes[i]);
      continue;
    }
    uint8_t Val = ShadowBytes[i];
    if (!AsanSetShadowFunc[Val]) {
      continue;
    }
    // Skip same values.
    while (j < End && ShadowMask[j] && Val == ShadowBytes[j]) {
      ++j;
    }

    if (j - i >= threshold) {
      copyToShadowInline(ShadowMask, ShadowBytes, Done, i, mirBuilder, ShadowBase, InsBefore);
      Done = j;
    }
  }

  copyToShadowInline(ShadowMask, ShadowBytes, Done, End, mirBuilder, ShadowBase, InsBefore);
}

void FunctionStackPoisoner::copyToShadowInline(const std::vector<uint8_t> ShadowMask, const std::vector<uint8_t> ShadowBytes,
                                               size_t Begin, size_t End, MIRBuilder *mirBuilder, BaseNode *ShadowBase,
                                               StmtNode *InsBefore) {
  if (Begin >= End) {
    return;
  }

  const size_t LargestStoreSizeInBytes = std::min<size_t>(sizeof(uint64_t), ASan.LongSize / 8);

  for (size_t i = Begin; i < End;) {
    if (!ShadowMask[i]) {
      assert(!ShadowBytes[i]);
      ++i;
      continue;
    }

    size_t StoreSizeInBytes = LargestStoreSizeInBytes;
    // Fit store size into the range.
    while (StoreSizeInBytes > End - i) {
      StoreSizeInBytes /= 2;
    }

    // Minimize store size by trimming trailing zeros.
    for (size_t j = StoreSizeInBytes - 1; j && !ShadowMask[i + j]; --j) {
      while (j <= StoreSizeInBytes / 2) {
        StoreSizeInBytes /= 2;
      }
    }

    uint64_t Val = 0;
    for (size_t j = 0; j < StoreSizeInBytes; j++) {
      Val |= uint64_t(ShadowBytes[i + j]) << (8 * j);
    }

    BinaryNode *Ptr = mirBuilder->CreateExprBinary(OP_add, *IntptrTy, ShadowBase,
                                                   mirBuilder->CreateIntConst(i, IntptrTy->GetPrimType()));
    PrimType primType;
    switch (StoreSizeInBytes * 8) {
      case 8:
        primType = PTY_i8;
        break;
      case 16:
        primType = PTY_i16;
        break;
      case 32:
        primType = PTY_i32;
        break;
      case 64:
        primType = PTY_i64;
        break;
      default: {
        primType = PTY_unknown;
      }
    }
    ConstvalNode *Poison = mirBuilder->CreateIntConst(Val, primType);
    MIRType *ptrType = GlobalTables::GetTypeTable().GetOrCreatePointerType(
        GlobalTables::GetTypeTable().GetPrimType(primType)->GetTypeIndex());
    IassignNode *iassignNode = mirBuilder->CreateStmtIassign(*ptrType, 0, Ptr, Poison);
    iassignNode->InsertAfterThis(*InsBefore);
    i += StoreSizeInBytes;
  }
}

void FunctionStackPoisoner::CleanStackShadowBytes(uint32 bytes, uint64 offset, MIRBuilder *mirBuilder,
                                                  BaseNode* ShadowBase, StmtNode *InsBefore) {
  PrimType ptype;
  switch (bytes)
  {
  case 8:
    ptype = PTY_u64;
    break;
  case 4:
    ptype = PTY_u32;
    break;
  case 2:
    ptype = PTY_u16;
    break;
  case 1:
    ptype = PTY_u8;
    break;
  default:
    CHECK_FATAL_FALSE("Not supported yet.");
  }
  BinaryNode *Ptr = mirBuilder->CreateExprBinary(OP_add, *IntptrTy, ShadowBase,
                                                 mirBuilder->CreateIntConst(offset, IntptrTy->GetPrimType()));
  MIRType* ptrType = GlobalTables::GetTypeTable().GetOrCreatePointerType(
        GlobalTables::GetTypeTable().GetPrimType(ptype)->GetTypeIndex());
  IassignNode *iassignNode = mirBuilder->CreateStmtIassign(*ptrType, 0, Ptr,
                                                           mirBuilder->CreateIntConst(0, PTY_u64));
  iassignNode->InsertAfterThis(*InsBefore);
}

void FunctionStackPoisoner::CleanStackShadow(const std::vector<uint8_t> ShadowMask, MIRBuilder *mirBuilder,
                                             BaseNode *ShadowBase, StmtNode *InsBefore) {
  const size_t LargestStoreSizeInBytes = std::min<size_t>(sizeof(uint64_t), ASan.LongSize / 8);
  size_t i = 0;
  bool non_zero = false;
  #define ASAN_CHECK_NON_ZERO(_i, _size) \
    non_zero = false; \
    for (size_t _j = 0; _j < (_size); _j++) { \
      if (ShadowMask[(_i) + _j]) { \
        non_zero = true; \
        break; \
      } \
    }

  for (; i + LargestStoreSizeInBytes <= ShadowMask.size();) {
    // we treat the last range specifically
    ASAN_CHECK_NON_ZERO(i, LargestStoreSizeInBytes);
    if (non_zero) {
      CleanStackShadowBytes(LargestStoreSizeInBytes, i, mirBuilder, ShadowBase, InsBefore);
    }
    i += LargestStoreSizeInBytes;
  }
  if (i + 4 <= ShadowMask.size()) {
    ASAN_CHECK_NON_ZERO(i, 4);
    if (non_zero) {
      CleanStackShadowBytes(4, i, mirBuilder, ShadowBase, InsBefore);
    }
    i += 4;
  }
  if (i + 2 <= ShadowMask.size()) {
    ASAN_CHECK_NON_ZERO(i, 2)
    if (non_zero) {
      CleanStackShadowBytes(2, i, mirBuilder, ShadowBase, InsBefore);
    }
    i += 2;
  }
  if (i + 1 <= ShadowMask.size()) {
    ASAN_CHECK_NON_ZERO(i, 1)
    if (non_zero) {
      CleanStackShadowBytes(1, i, mirBuilder, ShadowBase, InsBefore);
    }
    i += 1;
  }
  #undef ASAN_CHECK_NON_ZERO
}

void FunctionStackPoisoner::initializeCallbacks(const MIRModule &M) {
  MIRBuilder *mirBuilder = M.GetMIRBuilder();
#ifdef ENABLERBTREE
#else
  for (size_t Val : {0x00, 0xf1, 0xf2, 0xf3, 0xf5, 0xf8}) {
    std::ostringstream Name;
    Name << kAsanSetShadowPrefix;
    Name << std::setw(2) << std::setfill('0') << std::hex << Val;
    AsanSetShadowFunc[Val] = getOrInsertFunction(mirBuilder, Name.str().c_str(), GlobalTables::GetTypeTable().GetVoid(),
                                                 {IntptrTy, IntptrTy});
  }
#endif
  AsanAllocaPoisonFunc =
      getOrInsertFunction(mirBuilder, kAsanAllocaPoison, GlobalTables::GetTypeTable().GetVoid(), {IntptrTy, IntptrTy});
  AsanAllocasUnpoisonFunc = getOrInsertFunction(mirBuilder, kAsanAllocasUnpoison,
                                                GlobalTables::GetTypeTable().GetVoid(), {IntptrTy, IntptrTy});
}

std::set<MIRSymbol *> FunctionStackPoisoner::GetStackVarReferedByCallassigned() {
  std::set<MIRSymbol *> retOfCallassigned;
  MIRSymbolTable *symbolTable = mirFunction->GetSymTab();
  for (StmtNode &stmt : mirFunction->GetBody()->GetStmtNodes()) {
    if (stmt.GetOpCode() == OP_callassigned) {
      CallNode *tmp = dynamic_cast<CallNode *>(&stmt);
      const CallReturnVector *retVecPtr = tmp->GetCallReturnVector();
      // get the return value being refered by a function call
      for (const CallReturnPair &retPair : *retVecPtr) {
        MIRSymbol *symbol = symbolTable->GetSymbolFromStIdx(retPair.first.Idx());
        retOfCallassigned.insert(symbol);
      }
    }
  }
  return retOfCallassigned;
}

void FunctionStackPoisoner::collectLocalVariablesWithoutAlloca() {
  // Collect all variables refered by return statement of callassigned
  std::set<MIRSymbol *> retOfCallassigned = GetStackVarReferedByCallassigned();
  MIRSymbolTable *symbolTable = mirFunction->GetSymTab();
  size_t size = symbolTable->GetSymbolTableSize();
  CHECK_FATAL(size < UINT32_MAX, "Too large symbol table size.");
  for (uint32_t i = 0; i < uint32_t(size); ++i) {
    MIRSymbol *symbol = symbolTable->GetSymbolFromStIdx(LabelIdx(i));
    if (symbol == nullptr) {
      continue;
    }
    if (symbol->IsDeleted() || symbol->GetName() == "") {
      continue;
    }
    if (ASan.isInterestingSymbol(*symbol)) {
      if (StringUtils::StartsWith(symbol->GetName(), "asan_")) {
        continue;
      }
      // we skip symbols being refered by callassigned as a return value
      if (retOfCallassigned.find(symbol) != retOfCallassigned.end()) {
        continue;
      }
      StackAlignment = std::max(StackAlignment, symbol->GetType()->GetAlign());
      ASanStackVariableDescription description = {symbol->GetName(),
                                                  symbol->GetType()->GetSize(),
                                                  0,
                                                  symbol->GetType()->GetAlign(),
                                                  symbol,
                                                  nullptr,
                                                  0,
                                                  symbol->GetSrcPosition().LineNum()};
      stackVariableDesc.push_back(description);
    }
  }
}

void FunctionStackPoisoner::collectDescFromUnaryStmtNode(UnaryStmtNode &assignNode) {
  BaseNode *baseNode = assignNode.GetRHS();
  while (baseNode) {
    UnaryNode *rhs = dynamic_cast<UnaryNode *>(baseNode);
    if (rhs == nullptr) {
      return;
    }
    if (rhs->GetOpCode() == OP_alloca && ASan.isInterestingAlloca(*rhs)) {
      ConstvalNode *constvalNode = dynamic_cast<ConstvalNode *>(rhs->Opnd(0));
      if (constvalNode && isInFirstBlock(&assignNode)) {
        // static alloca
        MIRIntConst *mirConst = dynamic_cast<MIRIntConst *>(constvalNode->GetConstVal());
        ASanStackVariableDescription description = {
            "",      static_cast<size_t>(mirConst->GetValue().GetZXTValue()),
            0,       0,
            nullptr, &assignNode,
            0,       assignNode.GetSrcPos().LineNum()};
        stackVariableDesc.push_back(description);
      } else {
        // dynamic alloca
        ASanDynaVariableDescription description = {"", rhs->Opnd(0), &assignNode, 0, assignNode.GetSrcPos().LineNum()};
        dynamicAllocaDesc.push_back(description);
      }
    }
    baseNode = rhs->Opnd(0);
  }
}

void FunctionStackPoisoner::collectLocalVariablesWithAlloca() {
  for (StmtNode &stmt : mirFunction->GetBody()->GetStmtNodes()) {
    if (stmt.GetOpCode() == OP_regassign || stmt.GetOpCode() == OP_dassign) {
      UnaryStmtNode *assignNode = dynamic_cast<UnaryStmtNode *>(&stmt);
      CHECK_FATAL(assignNode != nullptr, "Node with OP_regassign or OP_dassign is not UnaryStmtNode.");
      collectDescFromUnaryStmtNode(*assignNode);
    }
  }
}

bool FunctionStackPoisoner::runOnFunction() {
  // Collect alloca, ret, etc.
  std::vector<StmtBlockNodePair> allStmtBlock = getAllOrderedStmtNodeList(mirFunction->GetBody());
  for (StmtBlockNodePair &stmtBlockPair : allStmtBlock) {
    StmtNode* stmtptr = stmtBlockPair.first;
    if (stmtptr->GetOpCode() == OP_return) {
      RetVec.push_back(stmtBlockPair);
    }
  }
  // Collect local variable
  // initialize stackVariableDesc and dynamicAllocaDesc
  collectLocalVariablesWithoutAlloca();
  collectLocalVariablesWithAlloca();
  // ignore variables already been used for alloca
  auto iter = stackVariableDesc.begin();
  while (iter != stackVariableDesc.end()) {
    if (iter->Symbol != nullptr && isUsedInAlloca[iter->Symbol]) {
      iter = stackVariableDesc.erase(iter);
    } else {
      ++iter;
    }
  }
  if (stackVariableDesc.empty() && dynamicAllocaDesc.empty()) {
    return false;
  }

  if (!dynamicAllocaDesc.empty()) {
    createDynamicAllocasInitStorage();
    for (auto &AI : dynamicAllocaDesc) {
      handleDynamicAllocaCall(&AI);
    }
    unpoisonDynamicAllocas();
  }
  processStackVariable();
  return true;
}

bool FunctionStackPoisoner::isInFirstBlock(StmtNode *stmtNode) {
  while (stmtNode) {
    if (stmtNode->IsCondBr()) {
      CondGotoNode *condGotoNode = dynamic_cast<CondGotoNode *>(stmtNode);
      ConstvalNode *constvalNode = dynamic_cast<ConstvalNode *>(condGotoNode->Opnd(0));
      if (constvalNode == nullptr) {
        // the stmt is after a direct branch stmt, cannot be the stmt in the first block
        return false;
      }
      MIRIntConst *mirIntConst = dynamic_cast<MIRIntConst *>(constvalNode->GetConstVal());
      if (!(mirIntConst && mirIntConst->GetValue() == 1)) {
        return false;
      }
    }
    stmtNode = stmtNode->GetPrev();
  }
  return true;
}

void FunctionStackPoisoner::createDynamicAllocasInitStorage() {
  MIRBuilder *mirBuilder = module->GetMIRBuilder();
  DynamicAllocaLayout = getOrCreateSymbol(mirBuilder, IntptrTy->GetTypeIndex(), "asan_dynamic_alloca", kStVar, kScAuto,
                                          mirFunction, kScopeLocal);
  DynamicAllocaLayout->GetAttrs().SetAlign(32);
  DassignNode *dassignNode = mirBuilder->CreateStmtDassign(DynamicAllocaLayout->GetStIdx(), 0,
                                                           mirBuilder->CreateIntConst(0, IntptrTy->GetPrimType()));
  mirFunction->GetBody()->InsertBefore(mirFunction->GetBody()->GetFirst(), dassignNode);
}

void FunctionStackPoisoner::unpoisonDynamicAllocasBeforeInst(StmtBlockNodePair stmtBlockPair) {
  MIRBuilder *mirBuilder = module->GetMIRBuilder();
  MapleVector<BaseNode *> args(mirBuilder->GetCurrentFuncCodeMpAllocator()->Adapter());
  args.emplace_back(mirBuilder->CreateDread(*DynamicAllocaLayout, IntptrTy->GetPrimType()));
  args.emplace_back(mirBuilder->CreateAddrof(*DynamicAllocaLayout, PTY_u64));

  CallNode *callNode = mirBuilder->CreateStmtCall(AsanAllocasUnpoisonFunc->GetPuidx(), args);
  BlockNode* block = stmtBlockPair.second;
  block->InsertBefore(stmtBlockPair.first, callNode);
}

// Unpoison dynamic allocas redzones.
void FunctionStackPoisoner::unpoisonDynamicAllocas() {
  for (auto &Ret : RetVec) {
    unpoisonDynamicAllocasBeforeInst(Ret);
  }
}

void FunctionStackPoisoner::handleDynamicAllocaCall(ASanDynaVariableDescription *AI) {
  MIRBuilder *mirBuilder = module->GetMIRBuilder();
  const unsigned Align = std::max(kAllocaRzSize, 1U);
  const uint64_t AllocaRedzoneMask = kAllocaRzSize - 1U;

  ConstvalNode *Zero = mirBuilder->CreateIntConst(0, IntptrTy->GetPrimType());
  ConstvalNode *AllocaRzSize = mirBuilder->CreateIntConst(kAllocaRzSize, IntptrTy->GetPrimType());
  ConstvalNode *AllocaRzMask = mirBuilder->CreateIntConst(AllocaRedzoneMask, IntptrTy->GetPrimType());

  BaseNode *NewSize;
  if (AI->Size->GetOpCode() == OP_constval) {
    ConstvalNode *constvalNode = dynamic_cast<ConstvalNode *>(AI->Size);
    MIRIntConst *intConst = dynamic_cast<MIRIntConst *>(constvalNode->GetConstVal());
    uint64_t PartialSize = uint64_t(intConst->GetValue().GetExtValue()) & AllocaRedzoneMask;
    uint64_t MisAlign = kAllocaRzSize - PartialSize;
    uint64_t PartialPadding = MisAlign;
    if (uint64_t(kAllocaRzSize) == MisAlign) {
      PartialPadding = 0;
    }
    uint64_t AdditionalChunkSize = Align + kAllocaRzSize + PartialPadding;
    NewSize =
        mirBuilder->CreateIntConst(AdditionalChunkSize + intConst->GetValue().GetExtValue(), IntptrTy->GetPrimType());
  } else {
    BaseNode *OldSize = AI->Size;
    BinaryNode *PartialSize = mirBuilder->CreateExprBinary(OP_band, *IntptrTy, OldSize, AllocaRzMask);
    BaseNode *Misalign = mirBuilder->CreateExprBinary(OP_sub, *IntptrTy, AllocaRzSize, PartialSize);

    MIRSymbol *misAlignSym = getOrCreateSymbol(mirBuilder, IntptrTy->GetTypeIndex(), "asan_misAlign", kStVar, kScAuto,
                                               mirFunction, kScopeLocal);
    DassignNode *dassignNode = mirBuilder->CreateStmtDassign(misAlignSym->GetStIdx(), 0, Misalign);
    dassignNode->InsertAfterThis(*AI->AllocaInst);

    Misalign = mirBuilder->CreateDread(*misAlignSym, IntptrTy->GetPrimType());
    BinaryNode *Cond = mirBuilder->CreateExprCompare(OP_ne, *IntptrTy, *IntptrTy, Misalign, AllocaRzSize);
    TernaryNode *PartialPadding = mirBuilder->CreateExprTernary(OP_select, *IntptrTy, Cond, Misalign, Zero);
    // Align is added to locate left redzone, PartialPadding for possible
    // partial redzone and kAllocaRzSize for right redzone respectively.
    BinaryNode *AdditionalChunkSize = mirBuilder->CreateExprBinary(
        OP_add, *IntptrTy, mirBuilder->CreateIntConst(Align + kAllocaRzSize, IntptrTy->GetPrimType()), PartialPadding);

    NewSize = mirBuilder->CreateExprBinary(OP_add, *IntptrTy, OldSize, AdditionalChunkSize);
  }
  // Insert new alloca with new NewSize and Align params.
  MIRType *ptrType = GlobalTables::GetTypeTable().GetOrCreatePointerType(*GlobalTables::GetTypeTable().GetInt8());
  MIRSymbol *tmpAlloca =
      getOrCreateSymbol(mirBuilder, ptrType->GetTypeIndex(), "asan_dyn_tmp", kStVar, kScAuto, mirFunction, kScopeLocal);
  UnaryNode *NewAlloca = mirBuilder->CreateExprUnary(OP_alloca, *GlobalTables::GetTypeTable().GetAddr64(), NewSize);

  DassignNode *dassignNode = mirBuilder->CreateStmtDassign(tmpAlloca->GetStIdx(), 0, NewAlloca);
  dassignNode->InsertAfterThis(*AI->AllocaInst);
  assert(AI->AllocaInst->Opnd(0)->GetOpCode() == OP_alloca);

  BinaryNode *NewAddress = mirBuilder->CreateExprBinary(OP_add, *GlobalTables::GetTypeTable().GetPrimType(PTY_u64),
                                                        mirBuilder->CreateDread(*tmpAlloca, PTY_a64),
                                                        mirBuilder->CreateIntConst(Align, PTY_u64));
  AI->AllocaInst->SetOpnd(NewAddress, 0);

  MapleVector<BaseNode *> args(mirBuilder->GetCurrentFuncCodeMpAllocator()->Adapter());
  args.emplace_back(NewAddress);
  args.emplace_back(AI->Size);
  CallNode *callNode = mirBuilder->CreateStmtCall(AsanAllocaPoisonFunc->GetPuidx(), args);
  callNode->InsertAfterThis(*AI->AllocaInst);

  // Insert __asan_alloca_poison call for new created alloca.
  dassignNode = mirBuilder->CreateStmtDassign(
      DynamicAllocaLayout->GetStIdx(), 0,
      mirBuilder->CreateExprTypeCvt(OP_cvt, PTY_i64, PTY_u64, *mirBuilder->CreateDread(*tmpAlloca, PTY_a64)));
  dassignNode->InsertAfterThis(*AI->AllocaInst);
}

MIRSymbol *FunctionStackPoisoner::createAllocaForLayout(StmtNode *insBefore, MIRBuilder *mirBuilder,
                                                        const ASanStackFrameLayout &L, MIRSymbol** asan_tmp_ptr) {
  CHECK_FATAL(L.FrameSize < UINT32_MAX, "Too large frame size.");
  MIRArrayType *arrayType =
      GlobalTables::GetTypeTable().GetOrCreateArrayType(*GlobalTables::GetTypeTable().GetInt8(), uint32_t(L.FrameSize));
  MIRSymbol *tmp =
      getOrCreateSymbol(mirBuilder, arrayType->GetTypeIndex(), "asan_tmp", kStVar, kScAuto, mirFunction, kScopeLocal);
  size_t realignStack = 32;
  assert((realignStack & (realignStack - 1)) == 0);
  size_t frameAlignment = std::max(L.FrameAlignment, realignStack);
  CHECK_FATAL(frameAlignment < UINT32_MAX, "Too large frameAlignment.");
  tmp->GetAttrs().SetAlign(uint32_t(frameAlignment));
  MIRType *ptrType = GlobalTables::GetTypeTable().GetOrCreatePointerType(*GlobalTables::GetTypeTable().GetInt8());
  MIRSymbol *alloca =
      getOrCreateSymbol(mirBuilder, ptrType->GetTypeIndex(), "asan_alloca", kStVar, kScAuto, mirFunction, kScopeLocal);
  DassignNode *dassignNode =
      mirBuilder->CreateStmtDassign(alloca->GetStIdx(), 0, mirBuilder->CreateAddrof(*tmp, PTY_u64));
  mirFunction->GetBody()->InsertBefore(insBefore, dassignNode);
  *asan_tmp_ptr = tmp;
  return alloca;
}

bool FunctionStackPoisoner::isFuncCallArg(const MIRSymbol *const symbolPtr) const {
  auto iter = callArgSymbols.find(symbolPtr);
  return iter != callArgSymbols.end();
}

bool FunctionStackPoisoner::isFuncCallArg(const std::string symbolName) const {
  auto iter = callArgSymbolNames.find(symbolName);
  return iter != callArgSymbolNames.end();
}

void FunctionStackPoisoner::processStackVariable() {
  if (stackVariableDesc.empty()) {
    return;
  }
  BlockNode* funcBody = mirFunction->GetBody();
  // we add an empty block at the head of the body
  // so that if the first statement is being replaced
  // we can still insert a instructions before the first
  BlockNode* insBefore = this->module->GetMemPool()->New<BlockNode>();
  funcBody->InsertFirst(insBefore);
  size_t granularity = 1ULL << Mapping.Scale;
  size_t minHeaderSize = std::max(ASan.LongSize / 2, granularity);
  const ASanStackFrameLayout &L = ComputeASanStackFrameLayout(stackVariableDesc, granularity, minHeaderSize);
  auto descriptionString = ComputeASanStackFrameDescription(stackVariableDesc);
  LogInfo::MapleLogger() << descriptionString << " --- " << L.FrameSize << "\n";

  bool doStackMalloc = true;
  uint64_t localStackSize = L.FrameSize;
  if (ENABLE_STACK_SIZE_LIMIT) {
    doStackMalloc = localStackSize <= kMaxStackMallocSize;
    int stackMallocIdx = StackMallocSizeClass(localStackSize);
    CHECK_FATAL(stackMallocIdx <= kMaxAsanStackMallocSizeClass, "Too large stackMallocIdx");
  }
  doStackMalloc = (!HasNonEmptyInlineAsm) && (!HasReturnsTwiceCall) && doStackMalloc;

  MIRBuilder *mirBuilder = module->GetMIRBuilder();
  MIRSymbol *asan_tmp;
  MIRSymbol *allocaValue = createAllocaForLayout(insBefore, mirBuilder, L, &asan_tmp);

  for (size_t i = 0; i < stackVariableDesc.size(); i++) {
    ASanStackVariableDescription desc = stackVariableDesc.at(i);
    if (desc.Symbol != nullptr) {
      MIRSymbol *localVar = desc.Symbol;
      if (localVar->IsStatic() || localVar->IsPUStatic()) {
        // skip them as I do not know how to assign its initial values
        continue;
      }
      MIRConst* localVarConst = localVar->GetKonst();
      if (localVarConst != nullptr) {
        MIRType& type = localVarConst->GetType();
        if (type.IsMIRArrayType()) {
          // Given an array type, we do not replace it
          // following accesses should be indirect accesses,
          // and they are interesting memory to be instrumented
          continue;
        }
      }
      BinaryNode *addExpr =
          mirBuilder->CreateExprBinary(OP_add, *IntptrTy,
                                       mirBuilder->CreateExprTypeCvt(OP_cvt, IntptrTy->GetPrimType(), PTY_u64,
                                                                     *mirBuilder->CreateDread(*allocaValue, PTY_a64)),
                                       mirBuilder->CreateIntConst(desc.Offset, IntptrTy->GetPrimType()));
      // change the variable to a pointer with name asan_<Var_name>
      MIRType *localVarPtr = GlobalTables::GetTypeTable().GetOrCreatePointerType(desc.Symbol->GetTyIdx());
      MIRSymbol *newLocalVar = getOrCreateSymbol(mirBuilder, localVarPtr->GetTypeIndex(), "asan_" + localVar->GetName(),
                                                 kStVar, kScAuto, mirFunction, kScopeLocal);
      newLocalVar->SetSrcPosition(localVar->GetSrcPosition());
      // initialize the field of the Var by dassign
      DassignNode *dassignNode = mirBuilder->CreateStmtDassign(newLocalVar->GetStIdx(), 0, addExpr);
      funcBody->InsertBefore(insBefore, dassignNode);
      // replace the Var being referenced
      replaceAllUsesWith(localVar, newLocalVar);
      /* The stack Var could be a function call parameter
      When the stack Var is initialized by the calling convention, its shadow memory can be
      initialized already, or it has value so that we need to assign the source value to its field
      NOTE: this must be done after replaceAllUsesWith; otherwise, the initialization could be falsely
      replaced*/
      if (isFuncCallArg(localVar)) {
        // MapleVector<BaseNode *> args(module->GetMPAllocator().Adapter());
        // args.emplace_back(mirBuilder->CreateDread(*newLocalVar, PTY_a64));
        // args.emplace_back(mirBuilder->CreateAddrof(*localVar, PTY_u64));
        // args.emplace_back(mirBuilder->GetConstUInt64(localVar->GetType()->GetSize()));
        // IntrinsiccallNode *intrinsiccallNode = mirBuilder->CreateStmtIntrinsicCall(INTRN_C_memcpy, args);
        // funcBody->InsertBefore(insBefore, intrinsiccallNode);
        AddrofNode* addrofNode = mirBuilder->CreateAddrof(*localVar, PTY_u64);
        DassignNode* dassignnode = mirBuilder->CreateStmtDassign(*newLocalVar, 0, addrofNode);
        funcBody->InsertBefore(insBefore, dassignnode);
      }
    }
    if (desc.AllocaInst != nullptr) {
      BinaryNode *addExpr =
          mirBuilder->CreateExprBinary(OP_add, *IntptrTy,
                                       mirBuilder->CreateExprTypeCvt(OP_cvt, IntptrTy->GetPrimType(), PTY_u64,
                                                                     *mirBuilder->CreateDread(*allocaValue, PTY_a64)),
                                       mirBuilder->CreateIntConst(desc.Offset, IntptrTy->GetPrimType()));
      if (desc.AllocaInst->GetOpCode() == OP_regassign || desc.AllocaInst->GetOpCode() == OP_dassign) {
        UnaryStmtNode *assignNode = dynamic_cast<UnaryStmtNode *>(desc.AllocaInst);
        assignNode->SetRHS(addExpr);
      }
    }
  }

  // The left-most redzone has enough space for at least 4 pointers.
  // Write the Magic value to redzone[0].
  BaseNode *basePlus0 = mirBuilder->CreateDread(*allocaValue, PTY_a64);
  IassignNode *basePlus0Store = mirBuilder->CreateStmtIassign(
      *IntptrPtrTy, 0, basePlus0, mirBuilder->CreateIntConst(kCurrentStackFrameMagic, IntptrTy->GetPrimType()));
  funcBody->InsertBefore(insBefore, basePlus0Store);
  // Write the frame description constant to redzone[1]
  BaseNode *basePlus1 = mirBuilder->CreateExprBinary(OP_add, *GlobalTables::GetTypeTable().GetPrimType(PTY_u64),
                                                     mirBuilder->CreateDread(*allocaValue, PTY_a64),
                                                     mirBuilder->CreateIntConst(ASan.LongSize / 8, PTY_u64));

  ConststrNode *description = module->CurFuncCodeMemPool()->New<ConststrNode>(
      PTY_a64, GlobalTables::GetUStrTable().GetOrCreateStrIdxFromName(descriptionString));

  IassignNode *basePlus1Store = mirBuilder->CreateStmtIassign(*IntptrPtrTy, 0, basePlus1, description);
  funcBody->InsertBefore(insBefore, basePlus1Store);
  // Write the PC to redzone[2]
  BaseNode *basePlus2 = mirBuilder->CreateExprBinary(OP_add, *GlobalTables::GetTypeTable().GetPrimType(PTY_u64),
                                                     mirBuilder->CreateDread(*allocaValue, PTY_a64),
                                                     mirBuilder->CreateIntConst(2 * ASan.LongSize / 8, PTY_u64));
  AddroffuncNode *addroffuncNode = mirBuilder->CreateExprAddroffunc(mirFunction->GetPuidx());
  addroffuncNode->SetPrimType(PTY_a64);
  IassignNode *basePlus2Store = mirBuilder->CreateStmtIassign(*IntptrPtrTy, 0, basePlus2, addroffuncNode);
  funcBody->InsertBefore(insBefore, basePlus2Store);
  const auto &shadowAfterScope = GetShadowBytesAfterScope(stackVariableDesc, L);

#ifdef ENABLERBTREE
  MapleVector<BaseNode *> args(mirBuilder->GetCurrentFuncCodeMpAllocator()->Adapter());
  args.emplace_back(mirBuilder->CreateDread(*allocaValue, PTY_a64));
  args.emplace_back(mirBuilder->CreateIntConst(shadowAfterScope.size() * L.Granularity, ASan.IntPtrPrim));
  auto callNode = mirBuilder->CreateStmtCall(ASan.AsanRBTStackInsert->GetPuidx(), args);
  callNode->InsertAfterThis(*insBefore);

  // Dig holes in redzone for variables
  for (auto const &desc : stackVariableDesc) {
    // LogInfo::MapleLogger() << "digging hole for " << desc.Name << "\n";
    MapleVector<BaseNode *> args(mirBuilder->GetCurrentFuncCodeMpAllocator()->Adapter());
    auto redzonePtr =
        mirBuilder->CreateExprBinary(OP_add, *ASan.IntPtrTy, mirBuilder->CreateDread(*allocaValue, PTY_a64),
                                     mirBuilder->CreateIntConst(desc.Offset, ASan.IntPtrPrim));
    args.emplace_back(redzonePtr);
    args.emplace_back(mirBuilder->CreateIntConst(desc.Size, ASan.IntPtrPrim));
    auto callNode = mirBuilder->CreateStmtCall(ASan.AsanRBTStackDelete->GetPuidx(), args);
    callNode->InsertAfterThis(*insBefore);
  }

  // (Un)poison the stack before all ret instructions.
  for (StmtNode *ret : RetVec) {
    MapleVector<BaseNode *> args(mirBuilder->GetCurrentFuncCodeMpAllocator()->Adapter());
    args.emplace_back(mirBuilder->CreateDread(*allocaValue, PTY_a64));
    args.emplace_back(mirBuilder->CreateIntConst(L.FrameSize, ASan.IntPtrPrim));
    auto callNode = mirBuilder->CreateStmtCall(ASan.AsanRBTStackDelete->GetPuidx(), args);
    callNode->InsertAfterThis(*ret);
  }
#else
  // Get the value of shadow memory
  MIRSymbol *shadowBase = getOrCreateSymbol(mirBuilder, IntptrTy->GetTypeIndex(), "asan_shadowBase", kStVar, kScAuto,
                                            mirFunction, kScopeLocal);
  DassignNode *dassignNode = mirBuilder->CreateStmtDassign(
      *shadowBase, 0, ASan.memToShadow(mirBuilder->CreateDread(*allocaValue, PTY_a64), *mirBuilder));
  funcBody->InsertBefore(insBefore, dassignNode);
  copyToShadow(shadowAfterScope, shadowAfterScope, mirBuilder,
               mirBuilder->CreateDread(*shadowBase, shadowBase->GetType()->GetPrimType()), insBefore);

  // (Un)poison the stack before all ret instructions.
  for (StmtBlockNodePair &ret : RetVec) {
    // Mark the current frame as retired.
    IassignNode *retiredNode = mirBuilder->CreateStmtIassign(
        *IntptrPtrTy, 0, basePlus0, mirBuilder->CreateIntConst(kRetiredStackFrameMagic, IntptrTy->GetPrimType()));
    ret.second->InsertBefore(ret.first, retiredNode);
    if (doStackMalloc) {
      CleanStackShadow(shadowAfterScope, mirBuilder,
                       mirBuilder->CreateDread(*shadowBase, shadowBase->GetType()->GetPrimType()), ret.first);
      // The following code is used for ensuring the stack is cleaned
      // BaseNode* tmpStackBase = mirBuilder->CreateAddrof(*asan_tmp, PTY_u64);
      // BaseNode* tmpSizeNode = mirBuilder->CreateIntConst(L.FrameSize, PTY_u64);
      // CallNode* call_stack_check = CreateStackCheck(mirBuilder, tmpStackBase, tmpSizeNode);
      // ret.second->InsertBefore(ret.first, call_stack_check);
    }
  }
#endif
  // We are done. Remove the old unused alloca instructions.
  for (ASanStackVariableDescription svd : stackVariableDesc) {
    if (svd.Symbol != nullptr) {
      svd.Symbol->SetIsDeleted();
    }
  }
  // remove the emptyBlock for insertion
  funcBody->RemoveStmt(insBefore);
}

BaseNode* FunctionStackPoisoner::TransformAddrofNode(MIRSymbol *oldVar, MIRSymbol *newVar, AddrofNode *addrofNode) {
  BaseNode *retNode = nullptr;
  MIRSymbol *mirSymbol = mirFunction->GetLocalOrGlobalSymbol(addrofNode->GetStIdx());
  if (mirSymbol->GetStIdx() == oldVar->GetStIdx()) {
    FieldID fieldID = addrofNode->GetFieldID();
    if (fieldID == 0) {
      retNode = module->GetMIRBuilder()->CreateDread(*newVar, PTY_a64);
    } else {
      // addrof a structure field
      auto type = dynamic_cast<MIRStructType*>(oldVar->GetType());
      OffsetPair offsetPair = type->GetFieldOffsetFromBaseAddr(fieldID);
      retNode = module->GetMIRBuilder()->CreateExprBinary(OP_add,
          *GlobalTables::GetTypeTable().GetPrimType(PTY_u64),
          module->GetMIRBuilder()->CreateDread(*newVar, PTY_a64),
          module->GetMIRBuilder()->CreateIntConst(offsetPair.byteOffset, PTY_u64));
    }
  }
  return retNode;
}

BaseNode* FunctionStackPoisoner::TransformDassignNode(MIRSymbol *oldVar, MIRSymbol *newVar, DassignNode *dassignNode) {
  BaseNode* retNode = nullptr;
  MIRSymbol *mirSymbol = mirFunction->GetLocalOrGlobalSymbol(dassignNode->GetStIdx());
  if (mirSymbol->GetStIdx() == oldVar->GetStIdx()) {
    BaseNode *newRHS = GetTransformedNode(oldVar, newVar, dassignNode->GetRHS());
    StmtNode *newStmtNode =
        module->GetMIRBuilder()->CreateStmtIassign(*newVar->GetType(), dassignNode->GetFieldID(),
                                                   module->GetMIRBuilder()->CreateDread(*newVar, PTY_a64), newRHS);
    retNode = newStmtNode;
  }
  return retNode;
}

BaseNode* FunctionStackPoisoner::TransformDreadNode(MIRSymbol *oldVar, MIRSymbol *newVar, DreadNode *dreadNode) {
  BaseNode* retNode = nullptr;
  MIRSymbol *mirSymbol = mirFunction->GetLocalOrGlobalSymbol(dreadNode->GetStIdx());
  if (mirSymbol->GetStIdx() == oldVar->GetStIdx()) {
    IreadNode *newStmtNode = module->GetMIRBuilder()->CreateExprIread(
        *GlobalTables::GetTypeTable().GetPrimType(dreadNode->GetPrimType()), *newVar->GetType(),
        dreadNode->GetFieldID(), module->GetMIRBuilder()->CreateDread(*newVar, PTY_a64));
    retNode = newStmtNode;
  }
  return retNode;
}

BaseNode* FunctionStackPoisoner::TransformDassignoffNode(MIRSymbol *oldVar, MIRSymbol *newVar, DassignoffNode *dassignoffNode) {
  BaseNode* retNode = nullptr;
  MIRSymbol *mirSymbol = mirFunction->GetLocalOrGlobalSymbol(dassignoffNode->GetStIdx());
  if (mirSymbol->GetStIdx() == oldVar->GetStIdx()) {
    BaseNode *newRHS = GetTransformedNode(oldVar, newVar, dassignoffNode->GetRHS());
    StmtNode *newStmtNode =
        module->GetMIRBuilder()->CreateStmtIassignoff(newVar->GetType()->GetPrimType(),
                                                      dassignoffNode->GetOffset(),
                                                      module->GetMIRBuilder()->CreateDread(*newVar, PTY_a64),
                                                      newRHS);
    retNode = newStmtNode;
  }
  return retNode;
}

BaseNode *FunctionStackPoisoner::GetTransformedNode(MIRSymbol *oldVar, MIRSymbol *newVar, BaseNode *baseNode) {
  BaseNode *retNode = nullptr;
  switch (baseNode->GetOpCode()) {
    case OP_addrof: {
      AddrofNode *addrofNode = getAddrofNode(baseNode);
      retNode = TransformAddrofNode(oldVar, newVar, addrofNode);
      if (retNode != nullptr) return retNode;
      break;
    }
    case OP_dassign: {
      DassignNode *dassignNode = dynamic_cast<DassignNode *>(baseNode);
      retNode = TransformDassignNode(oldVar, newVar, dassignNode);
      if (retNode != nullptr) return retNode;
      break;
    }
    case OP_dread: {
      DreadNode *dreadNode = getDreadNode(baseNode);
      retNode = TransformDreadNode(oldVar, newVar, dreadNode);
      if (retNode != nullptr) return retNode;
      break;
    }
    case OP_dassignoff: {
      DassignoffNode *dassignoffNode = dynamic_cast<DassignoffNode *>(baseNode);
      retNode = TransformDassignoffNode(oldVar, newVar, dassignoffNode);
      if (retNode != nullptr) return retNode;
      break;
    }
    case OP_block: {
      BlockNode* blockNode = dynamic_cast<BlockNode*>(baseNode);
      return this->replaceAllUsesOfBlockNode(oldVar, newVar, blockNode);
    }
    default: {
      break;
    }
  }
  
  // all other process must run following code to make sure every child has been visited
  for (size_t j = 0; j < baseNode->NumOpnds(); j++) {
    BaseNode *tmpNode = GetTransformedNode(oldVar, newVar, baseNode->Opnd(j));
    if (tmpNode != baseNode->Opnd(j)) {
      baseNode->SetOpnd(tmpNode, j);
    }
  }
  retNode = baseNode;
  return retNode;
}

BlockNode* FunctionStackPoisoner::replaceAllUsesOfBlockNode(MIRSymbol *oldVar, MIRSymbol *newVar, BlockNode* block) {
  // block->AddStatement(StmtNode*) has side effect on the input StmtNode
  // the GetPrev and GetNext will be changed automatically
  // Hence, we cache the node to be replaced, then call ReplaceStmt1WithStmt2 to replace them
  std::vector<std::pair<StmtNode *, StmtNode *>> toReplace;
  for (StmtNode &stmt1 : block->GetStmtNodes()) {
    StmtNode* stmt1ptr = &stmt1;
    BaseNode *newStmt = GetTransformedNode(oldVar, newVar, stmt1ptr);
    if (newStmt != dynamic_cast<BaseNode *>(stmt1ptr)) {
      StmtNode *stmt2ptr = dynamic_cast<StmtNode *>(newStmt);
      CHECK_FATAL(stmt2ptr != nullptr, "Get a stmt2 without StmtNode type");
      stmt2ptr->SetSrcPos(stmt1ptr->GetSrcPos());
      toReplace.push_back(std::pair<StmtNode*, StmtNode*>(stmt1ptr, stmt2ptr));
    }
  }

  for (auto ss : toReplace) {
    block->ReplaceStmt1WithStmt2(ss.first, ss.second);
  }
  return block;
}

void FunctionStackPoisoner::replaceAllUsesWith(MIRSymbol *oldVar, MIRSymbol *newVar) {
  if (mirFunction->GetBody() == nullptr) {
    return;
  }
  CHECK_FATAL(oldVar->GetTyIdx() == dynamic_cast<MIRPtrType *>(newVar->GetType())->GetPointedTyIdx(),
              "Replace Var SYmbol with different PointedTyIdx");
  this->replaceAllUsesOfBlockNode(oldVar, newVar, mirFunction->GetBody());
}
}  // namespace maple

#endif